{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Module 2 - Create Knowledge Base and Ingest Documents\n",
    "\n",
    "----\n",
    "\n",
    "This notebook provides sample code with step by step instructions for setting up an Amazon Bedrock Knowledge Base.\n",
    "\n",
    "----\n",
    "\n",
    "### Contents\n",
    "\n",
    "1. *Setup*\n",
    "1. *Create an S3 Data Source*\n",
    "1. *Setup AOSS Vector Index and Configure BKB Access Permissions*\n",
    "2. *Configure Amazon Bedrock Knowledge Base and Synchronize it with Data Source*\n",
    "3. *Conclusions and Next Steps*\n",
    "\n",
    "----\n",
    "\n",
    "### Introduction\n",
    "\n",
    "Foundation models (FMs) are powerful AI models trained on vast amounts of general-purpose data. However, many real-world applications require these models to generate responses grounded in domain-specific or proprietary information. Retrieval Augmented Generation (RAG) is a technique that enhances generative AI responses by retrieving relevant information from external data sources at query time.\n",
    "\n",
    "Amazon Bedrock Knowledge Bases (BKBs) provide a fully managed capability to implement RAG-based solutions. By integrating your own data — such as documents, manuals, and other domain-specific sources of information — into a knowledge base, you can improve the accuracy, relevance, and usefulness of model-generated responses. When a user submits a query, Amazon Bedrock Knowledge Bases search across the available data sources, retrieve the most relevant content, and pass this information to the foundation model to generate a more informed response.\n",
    "\n",
    "![BKB illustration](./images/bkb_illustration.png)\n",
    "\n",
    "This notebook demonstrates how to create an empty Amazon OpenSearch Serverless (AOSS) index, build an Amazon Bedrock Knowledge Base, and ingest documents into it for retrieval-augmented generation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Pre-requisites\n",
    "\n",
    "Please make sure that you have enabled the following model access in _Amazon Bedrock Console_:\n",
    "- `Amazon Titan Text Embeddings V2`.\n",
    "\n",
    "**If you are running AWS-facilitated event**, all other pre-requisites are satisfied and you can go to the next section.\n",
    "\n",
    "**If you are running this notebook as a self-paced lab**, then please note that this notebook requires permissions to:\n",
    "- create and delete *Amazon IAM* roles\n",
    "- create, update and delete *Amazon S3* buckets\n",
    "- access to *Amazon Bedrock*\n",
    "- access to *Amazon OpenSearch Serverless*\n",
    "\n",
    "In particular, if running on *SageMaker Studio*, you should add the following managed policies to your SageMaker execution role:\n",
    "- `IAMFullAccess`,\n",
    "- `AWSLambda_FullAccess`,\n",
    "- `AmazonS3FullAccess`,\n",
    "- `AmazonBedrockFullAccess`,\n",
    "- Custom policy for Amazon OpenSearch Serverless such as:\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"aoss:*\",\n",
    "            \"Resource\": \"*\"\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "````\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup\n",
    "\n",
    "### 1.1 Install and import the required libraries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --force-reinstall --no-deps -r ./requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Restart kernel\n",
    "from IPython.core.display import HTML\n",
    "HTML(\"<script>Jupyter.notebook.kernel.restart()</script>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Standard library imports\n",
    "import os\n",
    "import sys\n",
    "import json\n",
    "import time\n",
    "import random\n",
    "\n",
    "# Third-party imports\n",
    "import boto3\n",
    "from botocore.exceptions import ClientError\n",
    "\n",
    "# Local imports\n",
    "import utility\n",
    "\n",
    "# Print SDK versions\n",
    "print(f\"Python version: {sys.version.split()[0]}\")\n",
    "print(f\"Boto3 SDK version: {boto3.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Initial setup for clients and global variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create boto3 session and set AWS region\n",
    "boto_session = boto3.Session()\n",
    "aws_region = boto_session.region_name\n",
    "\n",
    "# Create boto3 clients for AOSS, Bedrock, and S3 services\n",
    "aoss_client = boto3.client('opensearchserverless')\n",
    "bedrock_agent_client = boto3.client('bedrock-agent')\n",
    "s3_client = boto3.client('s3')\n",
    "\n",
    "# Define names for AOSS, Bedrock, and S3 resources\n",
    "resource_suffix = random.randrange(100, 99999)\n",
    "s3_bucket_name = f\"bedrock-kb-{aws_region}-{resource_suffix}\"\n",
    "aoss_collection_name = f\"bedrock-kb-collection-{resource_suffix}\"\n",
    "aoss_index_name = f\"bedrock-kb-index-{resource_suffix}\"\n",
    "bedrock_kb_name = f\"bedrock-kb-{resource_suffix}\"\n",
    "\n",
    "# Set the Bedrock model to use for embedding generation\n",
    "embedding_model_id = 'amazon.titan-embed-text-v2:0'\n",
    "embedding_model_arn = f'arn:aws:bedrock:{aws_region}::foundation-model/{embedding_model_id}'\n",
    "embedding_model_dim = 1024\n",
    "\n",
    "# Some temporary local paths\n",
    "local_data_dir = 'data'\n",
    "\n",
    "# Print configurations\n",
    "print(\"AWS Region:\", aws_region)\n",
    "print(\"S3 Bucket:\", s3_bucket_name)\n",
    "print(\"AOSS Collection Name:\", aoss_collection_name)\n",
    "print(\"Bedrock Knowledge Base Name:\", bedrock_kb_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Create an S3 Data Source\n",
    "\n",
    "Amazon Bedrock Knowledge Bases can connect to a variety of data sources for downstream RAG applications. Supported data sources include Amazon S3, Confluence, Microsoft SharePoint, Salesforce, Web Crawler, and custom data sources.\n",
    "\n",
    "In this workshop, we will use Amazon S3 to store unstructured data — specifically, PDF files containing Amazon Shareholder Letters from different years. This S3 bucket will serve as the source of documents for our Knowledge Base. During the ingestion process, Bedrock will parse these documents, convert them into vector embeddings using an embedding model, and store them in a vector database for efficient retrieval during queries.\n",
    "\n",
    "### 2.1 Create an S3 bucket, if needed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if bucket exists, and if not create S3 bucket for KB data source\n",
    "\n",
    "try:\n",
    "    s3_client.head_bucket(Bucket=s3_bucket_name)\n",
    "    print(f\"Bucket '{s3_bucket_name}' already exists..\")\n",
    "except ClientError as e:\n",
    "    print(f\"Creating bucket: '{s3_bucket_name}'..\")\n",
    "    if aws_region == 'us-east-1':\n",
    "        s3_client.create_bucket(Bucket=s3_bucket_name)\n",
    "    else:\n",
    "        s3_client.create_bucket(\n",
    "            Bucket=s3_bucket_name,\n",
    "            CreateBucketConfiguration={'LocationConstraint': aws_region}\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Download data and upload to S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from urllib.request import urlretrieve\n",
    "\n",
    "# URLs of shareholder letters to download\n",
    "urls = [\n",
    "    'https://s2.q4cdn.com/299287126/files/doc_financials/2023/ar/2022-Shareholder-Letter.pdf',\n",
    "    'https://s2.q4cdn.com/299287126/files/doc_financials/2022/ar/2021-Shareholder-Letter.pdf',\n",
    "    'https://s2.q4cdn.com/299287126/files/doc_financials/2021/ar/Amazon-2020-Shareholder-Letter-and-1997-Shareholder-Letter.pdf',\n",
    "    'https://s2.q4cdn.com/299287126/files/doc_financials/2020/ar/2019-Shareholder-Letter.pdf'\n",
    "]\n",
    "\n",
    "# Corresponding local file names\n",
    "filenames = [\n",
    "    'AMZN-2022-Shareholder-Letter.pdf',\n",
    "    'AMZN-2021-Shareholder-Letter.pdf',\n",
    "    'AMZN-2020-Shareholder-Letter.pdf',\n",
    "    'AMZN-2019-Shareholder-Letter.pdf'\n",
    "]\n",
    "\n",
    "# Create local staging directory if it doesn't exist\n",
    "os.makedirs(local_data_dir, exist_ok=True)\n",
    "\n",
    "# Download each file and print confirmation\n",
    "for url, filename in zip(urls, filenames):\n",
    "    file_path = os.path.join(local_data_dir, filename)\n",
    "    urlretrieve(url, file_path)\n",
    "    print(f\"Downloaded: '{filename}' to '{local_data_dir}'..\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for root, _, files in os.walk(local_data_dir):\n",
    "    for file in files:\n",
    "        full_path = os.path.join(root, file)\n",
    "        s3_client.upload_file(full_path, s3_bucket_name, file)\n",
    "        print(f\"Uploaded: '{file}' to 's3://{s3_bucket_name}'..\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 Setup AOSS Vector Index and Configure BKB Access Permissions\n",
    "\n",
    "In this section, we’ll create a vector index using Amazon OpenSearch Serverless (AOSS) and configure the necessary access permissions for the Bedrock Knowledge Base (BKB) that we’ll set up later. AOSS provides a fully managed, serverless solution for running vector search workloads at billion-vector scale. It automatically handles resource scaling and eliminates the need for cluster management, while delivering low-latency, millisecond response times with pay-per-use pricing.\n",
    "\n",
    "While this example uses AOSS, it’s worth noting that Bedrock Knowledge Bases also supports other popular vector stores, including Amazon Aurora PostgreSQL with pgvector, Pinecone, Redis Enterprise Cloud, and MongoDB, among others\n",
    "\n",
    "### 3.1 Create IAM Role with Necessary Permissions for Bedrock Knowledge Base\n",
    "\n",
    "Let's first create an IAM role with all the necessary policies and permissions to allow BKB to execute operations, such as invoking Bedrock FMs and reading data from an S3 bucket. We will use a helper function for this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bedrock_kb_execution_role = utility.create_bedrock_execution_role(bucket_name=s3_bucket_name)\n",
    "bedrock_kb_execution_role_arn = bedrock_kb_execution_role['Role']['Arn']\n",
    "\n",
    "print(\"Created KB execution role with ARN:\", bedrock_kb_execution_role_arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Create AOSS Policies and Vector Collection\n",
    "\n",
    "Next we need to create and attach three key policies for securing and managing access to the AOSS collection: an encryption policy, a network access policy, and a data access policy. These policies ensure proper encryption, network security, and the necessary permissions for creating, reading, updating, and deleting collection items and indexes. This step is essential for configuring the OpenSearch collection to interact with BKB securely and efficiently (you can read more about AOSS collections [here](https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless.html)). We will use another helper function for this.\n",
    "\n",
    "⚠️ **Note:** _in order to keep setup overhead at mininum, in this example we **allow public internet access** to the OpenSearch Serverless collection resource. However, for production environments we strongly suggest to leverage private connection between your VPC and Amazon OpenSearch Serverless resources via an VPC endpoint, as described [here](https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-vpc.html)._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create AOSS policies for the new vector collection\n",
    "aoss_encryption_policy, aoss_network_policy, aoss_access_policy = utility.create_policies_in_oss(\n",
    "    vector_store_name=aoss_collection_name,\n",
    "    aoss_client=aoss_client,\n",
    "    bedrock_kb_execution_role_arn=bedrock_kb_execution_role_arn)\n",
    "\n",
    "print(\"Created encryption policy with name:\", aoss_encryption_policy['securityPolicyDetail']['name'])\n",
    "print(\"Created network policy with name:\", aoss_network_policy['securityPolicyDetail']['name'])\n",
    "print(\"Created access policy with name:\", aoss_access_policy['accessPolicyDetail']['name'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With all the necessary policies in place, let's proceed to actually creating a new AOSS collection. Please note that this can take a **few minutes to complete**. While you wait, you may want to [explore the AOSS Console](https://console.aws.amazon.com/aos/home?#opensearch/collections), where you will see your AOSS collection being created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Request to create AOSS collection\n",
    "aoss_collection = aoss_client.create_collection(name=aoss_collection_name, type='VECTORSEARCH')\n",
    "\n",
    "# Wait until collection becomes active\n",
    "print(\"Waiting until AOSS collection becomes active: \", end='')\n",
    "while True:\n",
    "    response = aoss_client.list_collections(collectionFilters={'name': aoss_collection_name})\n",
    "    status = response['collectionSummaries'][0]['status']\n",
    "    if status in ('ACTIVE', 'FAILED'):\n",
    "        print(\" done.\")\n",
    "        break\n",
    "    print('█', end='', flush=True)\n",
    "    time.sleep(5)\n",
    "\n",
    "print(\"An AOSS collection created:\", json.dumps(response['collectionSummaries'], indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Grant BKB Access to AOSS Data\n",
    "\n",
    "In this step, we create a data access policy that grants BKB the necessary permissions to read from our AOSS collections. We then attach this policy to the Bedrock execution role we created earlier, allowing BKB to securely access AOSS data when generating responses. We will be using helper function once again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aoss_policy_arn = utility.create_oss_policy_attach_bedrock_execution_role(\n",
    "    collection_id=aoss_collection['createCollectionDetail']['id'],\n",
    "    bedrock_kb_execution_role=bedrock_kb_execution_role)\n",
    "\n",
    "print(\"Waiting 60 sec for data access rules to be enforced: \", end='')\n",
    "for _ in range(12):  # 12 * 5 sec = 60 sec\n",
    "    print('█', end='', flush=True)\n",
    "    time.sleep(5)\n",
    "print(\" done.\")\n",
    "\n",
    "print(\"Created and attached policy with ARN:\", aoss_policy_arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Create an AOSS Vector Index\n",
    "\n",
    "Now that we have all necessary access permissions in place, we can create a vector index in the AOSS collection we created previously.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from requests_aws4auth import AWS4Auth\n",
    "from opensearchpy import OpenSearch, RequestsHttpConnection\n",
    "\n",
    "# Use default credential configuration for authentication\n",
    "credentials = boto_session.get_credentials()\n",
    "awsauth = AWS4Auth(\n",
    "    credentials.access_key,\n",
    "    credentials.secret_key,\n",
    "    aws_region,\n",
    "    'aoss',\n",
    "    session_token=credentials.token)\n",
    "\n",
    "# Construct AOSS endpoint host\n",
    "host = f\"{aoss_collection['createCollectionDetail']['id']}.{aws_region}.aoss.amazonaws.com\"\n",
    "\n",
    "# Build the OpenSearch client\n",
    "os_client = OpenSearch(\n",
    "    hosts=[{'host': host, 'port': 443}],\n",
    "    http_auth=awsauth,\n",
    "    use_ssl=True,\n",
    "    verify_certs=True,\n",
    "    connection_class=RequestsHttpConnection,\n",
    "    timeout=300\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to first define the index definiton with the desired indexing configuration, where we specify such things as number of shards and replicas of the index, vector embedding dimensions, the vector search engine (we are using FAISS here), as well as names and types of any other fields we need to have in the index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the configuration for the AOSS vector index\n",
    "index_definition = {\n",
    "   \"settings\": {\n",
    "      \"index.knn\": \"true\",\n",
    "       \"number_of_shards\": 1,\n",
    "       \"knn.algo_param.ef_search\": 512,\n",
    "       \"number_of_replicas\": 0,\n",
    "   },\n",
    "   \"mappings\": {\n",
    "      \"properties\": {\n",
    "         \"vector\": {\n",
    "            \"type\": \"knn_vector\",\n",
    "            \"dimension\": embedding_model_dim,\n",
    "             \"method\": {\n",
    "                 \"name\": \"hnsw\",\n",
    "                 \"engine\": \"faiss\",\n",
    "                 \"space_type\": \"l2\"\n",
    "             },\n",
    "         },\n",
    "         \"text\": {\n",
    "            \"type\": \"text\"\n",
    "         },\n",
    "         \"text-metadata\": {\n",
    "            \"type\": \"text\"\n",
    "         }\n",
    "      }\n",
    "   }\n",
    "}\n",
    "\n",
    "# Create an OpenSearch index\n",
    "response = os_client.indices.create(index=aoss_index_name, body=index_definition)\n",
    "\n",
    "# Waiting for index creation to propagate\n",
    "print(\"Waiting 30 sec for index update to propagate: \", end='')\n",
    "for _ in range(6):  # 6 * 5 sec = 30 sec\n",
    "    print('█', end='', flush=True)\n",
    "    time.sleep(5)\n",
    "print(\" done.\")\n",
    "\n",
    "print(\"A new AOSS index created:\", json.dumps(response, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Configure Amazon Bedrock Knowledge Base and Synchronize it with Data Source\n",
    "\n",
    "In this section, we’ll create an Amazon Bedrock Knowledge Base (BKB) and connect it to the data that will be stored in our newly created AOSS vector index.\n",
    "\n",
    "### 4.1 Create a Bedrock Knowledge Base\n",
    "\n",
    "Setting up a Knowledge Base involves providing two key configurations:\n",
    "1. **Storage Configuration** tells Bedrock where to store the generated vector embeddings by specifying the target vector store and providing the necessary connection detail (here, we use the AOSS vector index we created earlier),\n",
    "2. **Knowledge Base Configuration** defines how Bedrock should generate vector embeddings from your data by specifying the embedding model to use (`Titan Text Embeddings V2` in this sample), along with any additional settings required for handling multimodal content."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Vector Storage Configuration\n",
    "storage_config = {\n",
    "    \"type\": \"OPENSEARCH_SERVERLESS\",\n",
    "    \"opensearchServerlessConfiguration\": {\n",
    "        \"collectionArn\": aoss_collection[\"createCollectionDetail\"]['arn'],\n",
    "        \"vectorIndexName\": aoss_index_name,\n",
    "        \"fieldMapping\": {\n",
    "            \"vectorField\": \"vector\",\n",
    "            \"textField\": \"text\",\n",
    "            \"metadataField\": \"text-metadata\"\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "# Knowledge Base Configuration\n",
    "knowledge_base_config = {\n",
    "    \"type\": \"VECTOR\",\n",
    "    \"vectorKnowledgeBaseConfiguration\": {\n",
    "        \"embeddingModelArn\": embedding_model_arn\n",
    "    }\n",
    "}\n",
    "\n",
    "response = bedrock_agent_client.create_knowledge_base(\n",
    "    name=bedrock_kb_name,\n",
    "    description=\"Amazon shareholder letter knowledge base.\",\n",
    "    roleArn=bedrock_kb_execution_role_arn,\n",
    "    knowledgeBaseConfiguration=knowledge_base_config,\n",
    "    storageConfiguration=storage_config)\n",
    "\n",
    "bedrock_kb_id = response['knowledgeBase']['knowledgeBaseId']\n",
    "\n",
    "print(\"Waiting until BKB becomes active: \", end='')\n",
    "while True:\n",
    "    response = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=bedrock_kb_id)\n",
    "    if response['knowledgeBase']['status'] == 'ACTIVE':\n",
    "        print(\" done.\")\n",
    "        break\n",
    "    print('█', end='', flush=True)\n",
    "    time.sleep(5)\n",
    "\n",
    "print(\"A new Bedrock Knowledge Base created with ID:\", bedrock_kb_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's call a Bedrock API to get the information about our newly created Knowledge Base:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=bedrock_kb_id)\n",
    "\n",
    "print(json.dumps(response['knowledgeBase'], indent=2, default=str))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 Connect BKB to a Data Source\n",
    "\n",
    "With our Knowledge Base in place, the next step is to connect it to a data source. This involves two key actions:\n",
    "\n",
    "1. **Create a data source for the Knowledge Base** that will point to the location of our raw data (in this case, S3),\n",
    "2. **Define how that data should be processed and ingested into the vector store** — for example, by specifying a chunking configuration that controls how large each text fragment should be when generating vector embeddings for retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data Source Configuration\n",
    "data_source_config = {\n",
    "        \"type\": \"S3\",\n",
    "        \"s3Configuration\":{\n",
    "            \"bucketArn\": f\"arn:aws:s3:::{s3_bucket_name}\",\n",
    "            # \"inclusionPrefixes\":[\"*.*\"]   # you can use this if you want to create a KB using data within s3 prefixes.\n",
    "        }\n",
    "    }\n",
    "\n",
    "# Vector Ingestion Configuration\n",
    "vector_ingestion_config = {\n",
    "        \"chunkingConfiguration\": {\n",
    "            \"chunkingStrategy\": \"FIXED_SIZE\",\n",
    "            \"fixedSizeChunkingConfiguration\": {\n",
    "                \"maxTokens\": 512,\n",
    "                \"overlapPercentage\": 20\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "\n",
    "response = bedrock_agent_client.create_data_source(\n",
    "    name=bedrock_kb_name,\n",
    "    description=\"Amazon shareholder letter knowledge base.\",\n",
    "    knowledgeBaseId=bedrock_kb_id,\n",
    "    dataSourceConfiguration=data_source_config,\n",
    "    vectorIngestionConfiguration=vector_ingestion_config\n",
    ")\n",
    "\n",
    "bedrock_ds_id = response['dataSource']['dataSourceId']\n",
    "\n",
    "print(\"A new BKB data source created with ID:\", bedrock_ds_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also use Bedrock API to get the information about our newly created BKB data source:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "response = bedrock_agent_client.get_data_source(knowledgeBaseId=bedrock_kb_id, dataSourceId=bedrock_ds_id)\n",
    "\n",
    "print(json.dumps(response['dataSource'], indent=2, default=str))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3 Synchronize BKB with Data Source\n",
    "\n",
    "Once the Knowledge Base and its data source are configured, we can start a fully-managed data ingestion job. During this process, BKB will retrieve the documents from the connected data source (on S3, in this case), extract and preprocess the content, split it into smaller chunks based on the configured chunking strategy, generate vector embeddings for each chunk, and store those embeddings in the vector store (AOSS vector store, in this case).\n",
    "\n",
    "![BKB data ingestion](./images/data_ingestion.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Start an ingestion job\n",
    "response = bedrock_agent_client.start_ingestion_job(knowledgeBaseId=bedrock_kb_id, dataSourceId=bedrock_ds_id)\n",
    "\n",
    "bedrock_job_id = response['ingestionJob']['ingestionJobId']\n",
    "\n",
    "print(\"A new BKB ingestion job started with ID:\", bedrock_job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Wait until ingestion job completes\n",
    "print(\"Waiting until BKB ingestion job completes: \", end='')\n",
    "while True:\n",
    "    response = bedrock_agent_client.get_ingestion_job(\n",
    "        knowledgeBaseId = bedrock_kb_id,\n",
    "        dataSourceId = bedrock_ds_id,\n",
    "        ingestionJobId = bedrock_job_id)\n",
    "    if response['ingestionJob']['status'] == 'COMPLETE':\n",
    "        print(\" done.\")\n",
    "        break\n",
    "    print('█', end='', flush=True)\n",
    "    time.sleep(5)\n",
    "\n",
    "print(\"The BKB ingestion job finished:\", json.dumps(response['ingestionJob'], indent=2, default=str))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Conclusions and Next Steps\n",
    "\n",
    "In this notebook, we walked through the process of creating an Amazon Bedrock Knowledge Base (BKB) and ingesting documents to enable Retrieval Augmented Generation (RAG) capabilities. We started by setting up the environment, installing the required libraries, and initializing the necessary AWS service clients. Then, we created an Amazon S3 bucket to store unstructured data (PDF documents) and uploaded sample files. We proceeded by provisioning an Amazon OpenSearch Serverless (AOSS) collection and index, configuring the appropriate IAM roles and permissions, and granting access to the BKB. Finally, we created the BKB, connected it to the S3 data source, and synchronized the documents to generate vector embeddings, which were stored in AOSS."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next Steps\n",
    "\n",
    "Please execute next cell to store some important varables that will be needed in other notebooks of this module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store s3_bucket_name aoss_encryption_policy aoss_network_policy aoss_access_policy aoss_collection bedrock_kb_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now please go on to explore how you can interact with the newly created Knowledge Base via Bedrock APIs for RAG applications, please proceed to the next notebook:\n",
    "\n",
    "&nbsp; **NEXT ▶** [2_managed-rag-with-retrieve-and-generate-api.ipynb](./2\\_managed-rag-with-retrieve-and-generate-api.ipynb)."
   ]
  }
 ],
 "metadata": {
  "availableInstances": [
   {
    "_defaultOrder": 0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.t3.medium",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 1,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.t3.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 2,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.t3.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 3,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.t3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 4,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 5,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 6,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 7,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 8,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 9,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 10,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 11,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 12,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5d.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 13,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5d.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 14,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5d.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 15,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5d.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 16,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5d.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 17,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5d.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 18,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5d.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 19,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 20,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": true,
    "memoryGiB": 0,
    "name": "ml.geospatial.interactive",
    "supportedImageNames": [
     "sagemaker-geospatial-v1-0"
    ],
    "vcpuNum": 0
   },
   {
    "_defaultOrder": 21,
    "_isFastLaunch": true,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.c5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 22,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.c5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 23,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.c5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 24,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.c5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 25,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 72,
    "name": "ml.c5.9xlarge",
    "vcpuNum": 36
   },
   {
    "_defaultOrder": 26,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 96,
    "name": "ml.c5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 27,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 144,
    "name": "ml.c5.18xlarge",
    "vcpuNum": 72
   },
   {
    "_defaultOrder": 28,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.c5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 29,
    "_isFastLaunch": true,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g4dn.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 30,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g4dn.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 31,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g4dn.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 32,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g4dn.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 33,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g4dn.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 34,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g4dn.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 35,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 61,
    "name": "ml.p3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 36,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 244,
    "name": "ml.p3.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 37,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 488,
    "name": "ml.p3.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 38,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.p3dn.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 39,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.r5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 40,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.r5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 41,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.r5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 42,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.r5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 43,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.r5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 44,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.r5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 45,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.r5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 46,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.r5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 47,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 48,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 49,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 50,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 51,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 52,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 53,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.g5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 54,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.g5.48xlarge",
    "vcpuNum": 192
   },
   {
    "_defaultOrder": 55,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 56,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4de.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 57,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.trn1.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 58,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.trn1.32xlarge",
    "vcpuNum": 128
   },
   {
    "_defaultOrder": 59,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.trn1n.32xlarge",
    "vcpuNum": 128
   }
  ],
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
